package com.example.tree;

import java.util.ArrayList;
import java.util.List;

// 定义 Trie 节点类
class TrieNode {
    char value; // 节点的值（字符）
    boolean isEnd; // 是否是关键词的结束字符
    TrieNode[] children; // 孩子节点数组

    public TrieNode(char value) {
        this.value = value;
        this.isEnd = false;
        this.children = new TrieNode[26]; // 假设只处理小写字母
    }
}

// 定义 Trie 树类
class Trie {
    private TrieNode root;

    public Trie() {
        root = new TrieNode('/');
    }

    // 插入关键词
    public void insert(String word) {
        TrieNode node = root;
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            int index = c - 'a';
            if (node.children[index] == null) {
                node.children[index] = new TrieNode(c);
            }
            node = node.children[index];
        }
        node.isEnd = true;
    }

    // 查找前缀对应的关键词列表
    public List<String> findWordsWithPrefix(String prefix) {
        List<String> result = new ArrayList<>();
        TrieNode node = root;
        // 找到前缀的最后一个节点
        for (int i = 0; i < prefix.length(); i++) {
            char c = prefix.charAt(i);
            int index = c - 'a';
            if (node.children[index] == null) {
                return result; // 前缀不存在，直接返回空列表
            }
            node = node.children[index];
        }
        // 遍历以该节点为根的子树，收集所有关键词
        collectWords(node, prefix, result);
        return result;
    }

    // 递归遍历子树并收集关键词
    private void collectWords(TrieNode node, String prefix, List<String> result) {
        if (node.isEnd) {
            result.add(prefix);
        }
        for (int i = 0; i < node.children.length; i++) {
            if (node.children[i] != null) {
                collectWords(node.children[i], prefix + (char) ('a' + i), result);
            }
        }
    }

    public static void main(String[] args) {
        // 创建 Trie 树并插入关键词
        Trie trie = new Trie();
        String[] keywords = {"apple", "application", "banana", "bat", "batman", "batmobile", "cat", "car", "cartoon", "dog", "doghouse", "doggy"};
        for (String keyword : keywords) {
            trie.insert(keyword);
        }

        // 测试搜索关键词提示功能
        String inputPrefix = "ba";
        List<String> suggestions = trie.findWordsWithPrefix(inputPrefix);
        System.out.println("输入前缀 \"" + inputPrefix + "\" 的关键词提示:");
        for (String suggestion : suggestions) {
            System.out.println(suggestion);
        }
    }

}